package com.tgy.algorithm.base._图_01;

import java.util.*;
import java.util.stream.Collectors;

class UnionFindSet<V> {
    private HashMap<Node<V>,Node<V>> parents;
    private HashMap<Node<V>,Integer> sizeMap;

    UnionFindSet(List<Node<V>> nodes) {
        parents = new HashMap<>();
        sizeMap = new HashMap<>();
        for (Node<V> node : nodes) {
            parents.put(node,node);
            sizeMap.put(node,1);
        }
    }

    public Node<V> findParent(Node<V> node) {
        ArrayList<Node<V>> nodes = new ArrayList<>();
        while (node != parents.get(node)) {
            nodes.add(node);
            node = parents.get(node);
        }

        for (Node<V> nd : nodes) {
            parents.put(nd,node);
        }

        return node;
    }

    public boolean isSameSet(Node<V> one,Node<V> two) {
        return findParent(one) == findParent(two);
    }

    public void union(Node<V> one,Node<V> two) {

        Node<V> parentOne = findParent(one);
        Node<V> parentTwo = findParent(two);
        if (parentOne == parentTwo) {
            return;
        }

        Integer oneSize = sizeMap.get(parentOne);
        Integer twoSize = sizeMap.get(parentTwo);
        if (oneSize >= twoSize) {
             parents.put(parentTwo,parentOne);
             sizeMap.put(parentOne,oneSize + twoSize);
             sizeMap.remove(parentTwo);
        }else {
            parents.put(parentOne,parentTwo);
            sizeMap.put(parentTwo,oneSize + twoSize);
            sizeMap.remove(parentOne);
        }
    }
}

public class _最小生成树 {

    public static void heapTest() {

        PriorityQueue<Integer> integers = new PriorityQueue<>();
        int[] nums = {2,3,4,2,4,5,1};
        for (int i = 0; i < nums.length; i++) {
            integers.add(nums[i]);
        }

        while (!integers.isEmpty()) {
            System.out.println(integers.poll());
        }
    }

    public static List<Edge<Integer>> kruskal(Graph<Integer> graph) {

        ArrayList<Edge<Integer>> edges = new ArrayList<>();
        List<Node<Integer>> nodeList = graph.nodes.values().stream().collect(Collectors.toList());
        UnionFindSet<Integer> unionFindSet = new UnionFindSet<Integer>(nodeList);

        PriorityQueue<Edge<Integer>> minHeap = new PriorityQueue<>(new Comparator<Edge<Integer>>() {
            @Override
            public int compare(Edge<Integer> o1, Edge<Integer> o2) {
                return o1.value - o2.value;
            }
        });

        for (Edge<Integer> edge : graph.edges) {
            minHeap.offer(edge);
        }


        System.out.println(minHeap);

        while (!minHeap.isEmpty()) {
            Edge<Integer> edge = minHeap.poll();
            if (!unionFindSet.isSameSet(edge.to,edge.from)) {
                unionFindSet.union(edge.to,edge.from);
                edges.add(edge);
            }
        }

        return edges;
    }

    public static List<Edge<Integer>> prim(Graph<Integer> graph) {

        ArrayList<Edge<Integer>> res = new ArrayList<>();

        PriorityQueue<Edge<Integer>> minStack = new PriorityQueue<>(new Comparator<Edge<Integer>>() {
            @Override
            public int compare(Edge<Integer> o1, Edge<Integer> o2) {
                return o1.value - o2.value;
            }
        });

        HashSet<Node<Integer>> visitSet = new HashSet<>();
        Node<Integer> node = graph.nodes.get(2);
        visitSet.add(node);
        for (Edge<Integer> edge : node.edges) {
            minStack.offer(edge);
        }

        while (!minStack.isEmpty()) {
            Edge<Integer> visitEdge = minStack.poll();
            if (!visitSet.contains(visitEdge.to)) {
                visitSet.add(visitEdge.to);
                res.add(visitEdge);
                for (Edge<Integer> edge : visitEdge.to.edges) {
                    minStack.offer(edge);
                }
            }
        }

        return res;
    }

    public static void main(String[] args) {
//        heapTest();

        Graph<Integer> graph = _拓扑排序.createGraph(new int[][]{
                {1, 2, 10},
                {2, 1, 10},
                {1, 3, 2},
                {3, 1,  2},
                {1, 4, 2},
                {4, 1,  2},
                {1, 5, 3},
                { 5, 1, 3},
                {2, 3, 4},
                {3, 2, 4},
                {3, 4, 3},
                {4, 3,  3},
                {4, 5, 4},
                {5, 4,  4},
        });
        List<Edge<Integer>> kruskal = kruskal(graph);
        kruskal.sort((o1, o2) -> o1.value - o2.value);
        List<Edge<Integer>> prim = prim(graph);
        System.out.println(graph);
        System.out.println(kruskal);
//        System.out.println(prim.sort((o1, o2) -> o1.value - o2.value));
        prim.sort((o1, o2) -> o1.value - o2.value);
        System.out.println(prim);
    }
}
